# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import json
import logging
import os.path
import time
from fastapi import Request
from tortoise.transactions import in_transaction
from exception import AppException
from common.enums.apply import DrawsEnum
from common.enums.wallet import WalletEnum
from common.utils.urls import UrlUtil
from common.utils.times import TimeUtil
from common.utils.tools import ToolsUtil
from common.utils.config import ConfigUtil
from common.models.users import UserModel
from common.models.users import UserWalletModel
from common.models.draws import DrawsRecordsModel
from common.chain.draws_server import DrawsMjServer
from common.chain.draws_server import DrawsDoubaoServer
from common.service.subscribe_logic import SubscribeLogic
from plugins.censor.driver import CensorDriver
from plugins.storage.driver import StorageDriver
from apps.api.schemas.apply.draws import drawing_schema as schema

logger = logging.getLogger(__name__)


class DrawingService:

    def __init__(self, request: Request, user_id: int, params: schema.DrawGenerateIn):
        self.domain = request.scope["scheme"] + "://" + request.headers.get("host")
        self.request = request
        self.user_id = user_id
        self.params = params
        self.points = 0
        self.user = None
        self.type = DrawsEnum.TYPE_TEXT_TO_IMAGE
        self.vip_free = False

        # 参数数据
        self.prompt_eq = ""
        self.other_params = {}
        self.complex_params = {}

    async def generate(self):
        """
        文生图/图生图
        """
        # 验证参数
        await self._initialize()

        # 验证费用
        await self._check_charging()

        # 违规审查
        review = CensorDriver.REVIEW_DRAWS
        censor = await CensorDriver.text(user_id=self.user_id, text=self.params.prompt, scene=review)
        if censor is not None:
            raise AppException(f"描述词违规: {censor}")

        # 发起请求
        response = {}
        try:
            if self.params.engine == DrawsEnum.ENGINE_SD:
                pass
            elif self.params.engine == DrawsEnum.ENGINE_DALLE:
                salt = str(time.time()) + str(ToolsUtil.make_uuid())
                response["task_id"] = ToolsUtil.make_md5_str(salt)
                self.prompt_eq = self.params.prompt_in
                self.other_params = {
                    "quality": self.params.other_params.get("quality"),
                    "style": self.params.other_params.get("style")
                }
            elif self.params.engine == DrawsEnum.ENGINE_MJ:
                mj_server = DrawsMjServer({
                    "callback_url": self.domain + "/api/draws/drawing/notify_mj",
                    "action": self.params.action,
                    "model": self.params.model,
                    "scale": self.params.scale,
                    "version": self.params.version,
                    "image_id": self.params.image_id,
                    "image_ref": self.params.image_ref,
                    "prompt": self.prompt_eq,
                    "other_params": {
                        "mode": self.params.other_params.get("mode"),
                        "style": self.params.other_params.get("style"),
                        "no_content": self.params.other_params.get("no_content")
                    },
                    "complex_params": {
                        "seed": self.params.complex_params.get("seed") or -1,
                        "iw": int(self.params.complex_params.get("iw") or 1),
                        "q": float(self.params.complex_params.get("q") or 1),
                        "s": int(self.params.complex_params.get("s") or 100),
                        "c": int(self.params.complex_params.get("c") or 0)
                    }
                })
                response = await mj_server.imagine()
                self.prompt_eq = mj_server.prompt
                self.other_params = mj_server.other_params
                self.complex_params = mj_server.complex_params
            elif self.params.engine == DrawsEnum.ENGINE_DOUBAO:
                doubao_server = DrawsDoubaoServer({
                    "prompt": self.prompt_eq,
                    "model": self.params.model,
                    "scale": self.params.scale,
                    "image_ref": self.params.image_ref,
                    "complex_params": self.params.complex_params
                })
                response = await doubao_server.imagine()
                self.complex_params = doubao_server.complex_params
        except Exception as e:
            raise AppException(str(e))

        # 验证响应
        if not response or not response.get("task_id"):
            raise AppException("生成图片生成失败")

        # 保存记录
        await self._save_record(response.get("task_id"))

    async def _initialize(self):
        """ 初始化 """
        # 提示词
        self.prompt_eq = self.params.prompt_in
        if self.params.prompt_en:
            self.prompt_eq = self.params.prompt_en

        # UV操作
        if self.params.action != DrawsEnum.ACTION_GENERATE:
            if self.params.action.startswith(DrawsEnum.ACTION_UPSAMPLE):
                self.type = DrawsEnum.TYPE_UPSCALE_IMAGE
            elif self.params.action.startswith(DrawsEnum.ACTION_VARIATION):
                self.type = DrawsEnum.TYPE_VARIATION_IMAGE
            elif self.params.image_ref:
                self.type = DrawsEnum.TYPE_IMAGE_TO_IMAGE

        # 验证用户
        self.user = await UserModel.filter(id=self.user_id).first()
        if not self.user:
            raise AppException("用户信息错误")

    async def _check_charging(self):
        """ 验证费用 """
        # 验证通道状态
        config = await ConfigUtil.get("draws", self.params.engine) or {}
        if not config.get("status"):
            AppException("当前绘画渠道已关闭")

        # 获取所需积分
        if self.params.engine == DrawsEnum.ENGINE_MJ:
            mode = config.get("mode") or {}
            data = mode.get(self.params.other_params.get("mode"), {})
            self.points = int(data.get("points", "0"))
            if not data.get("status"):
                raise AppException("当前生成模式已停用")
        else:
            self.points = int(config.get("points", "0"))

        # 获取订阅会员
        member = await SubscribeLogic.get_surplus_free(user_id=self.user_id, scene="draws", mark=self.params.engine)
        if member:
            self.vip_free = member.get("free_status")

        # 验证账户余额
        if not self.vip_free and self.user.points < self.points:
            raise AppException("账户余额不足,请充值!")

    async def _save_record(self, task_id: str) -> None:
        """ 保存记录 """
        async with in_transaction("mysql"):
            if self.vip_free:
                self.points = 0
            else:
                # 用户扣费
                if self.points > 0:
                    self.user.points = self.user.points - self.points
                    await self.user.save()
                # 账户明细
                if self.points > 0:
                    await UserWalletModel.dec(
                        scene=WalletEnum.POINTS,
                        user_id=self.user_id,
                        change_type=WalletEnum.POINTS_DEC_DRAWS,
                        change_amount=self.points,
                        project="AI绘画"
                    )

            # 写入记录
            status = DrawsEnum.STATUS_ING if self.params.engine == "mj" else DrawsEnum.STATUS_WAIT
            code = ToolsUtil.make_md5_str(ToolsUtil.make_uuid() + task_id)
            await DrawsRecordsModel.create(
                engine=self.params.engine,
                code=code,
                type=self.type,
                user_id=self.user_id,
                action=self.params.action,
                model=self.params.model,
                scale=self.params.scale,
                version=self.params.version,
                prompt_in=self.params.prompt_in,
                prompt_en=self.params.prompt_en,
                prompt_eq=self.prompt_eq,
                other_params=json.dumps(self.params.other_params, ensure_ascii=False),
                complex_params=json.dumps(self.complex_params, ensure_ascii=False),
                task_id=task_id,
                points=self.points,
                is_free=int(self.vip_free),
                status=status,
                create_time=int(time.time()),
                update_time=int(time.time())
            )


class DrawsNotifyService:

    @classmethod
    async def notify_mj(cls, response: dict) -> None:
        """
        MJ回调
            {
                "image_url": "https://xx.cn/midjourney/45f6e5.png?imageMogr2/thumbnail/!50p",
                "image_width": 1024,
                "image_height": 1024,
                "actions": ["upscale1", "upscale4", "reroll", "variation1", "variation4"],
                "raw_image_url": "https://xx,cn/midjourney/45f6e5.png",
                "raw_image_width": 2048,
                "raw_image_height": 2048,
                "sub_image_urls": [],
                "progress": 100,
                "image_id": "1350021065006383104",
                "task_id": "bee5385b-9363-474a-88de-e7d16245f6e5",
                "success": true,
                "trace_id"?: "048c0b56-35d0-46c8-80ce-2dca766a7fbd"
            }

        Author:
            zero
        """
        code: str = response.get("code", "")
        task_id: str = response.get("task_id", "")

        album_dir_path = ""
        thumb_dir_path = ""
        record = None
        try:
            # 校验数据
            if not response.get("task_id"):
                return

            # 查询记录
            record = await DrawsRecordsModel.filter(task_id=task_id).first()
            if not record:
                return

            # 成功忽略
            if record.status == DrawsEnum.STATUS_YES and record.image_loc:
                return

            # 失败
            if code:
                error: str = response.get("detail", "-")
                await DrawsRecordsModel.filter(id=record.id).update(
                    status=DrawsEnum.STATUS_FAIL,
                    fail_reason=error,
                    notify_snap=json.dumps(response, ensure_ascii=False),
                    update_time=int(time.time())
                )

            # 成功
            else:
                image_id: str = response.get("image_id")
                image_url: str = response.get("raw_image_url")
                able_actions = response.get("actions", [])
                await DrawsRecordsModel.filter(id=record.id).update(
                    status=DrawsEnum.STATUS_YES,
                    image_id=image_id,
                    image_url=image_url,
                    thumbnail=image_url,
                    able_actions=json.dumps(able_actions),
                    notify_snap=json.dumps(response, ensure_ascii=False),
                    update_time=int(time.time())
                )

                # 名称信息
                name: str = os.path.basename(image_url.split("?")[0])
                ext = name.split(".")[-1].lower()
                if ext not in ["png", "jpg", "jpeg", "gif"]:
                    ext = "png"

                # 图片名称
                filename = ToolsUtil.make_md5_str(name) + "." + ext

                # 存储路径
                curr_now = TimeUtil.timestamp_to_date(int(time.time()), "%Y%m%d")
                album_dir: str = f"storage/drawing/albums/{curr_now}/{filename}"
                thumb_dir: str = f"storage/drawing/thumbnail/{curr_now}/{filename}"
                album_dir_path: str = UrlUtil.to_root_path(path=album_dir)
                thumb_dir_path: str = UrlUtil.to_root_path(path=thumb_dir)

                # 下载图片
                await ToolsUtil.download_file(
                    url=str(image_url),
                    save_path=album_dir_path
                )

                # 生成缩略图
                ToolsUtil.generate_thumbnail(
                    input_path=album_dir_path,
                    output_path=thumb_dir_path
                )

                # 图上传OSS
                engine: str = await UrlUtil.get_storage_engine()
                if engine != "local":
                    await StorageDriver.push(path=album_dir_path, key=album_dir)
                    await StorageDriver.push(path=thumb_dir_path, key=thumb_dir)

                # 更新图片
                await DrawsRecordsModel.filter(id=record.id).update(
                    image_loc=album_dir,
                    thumbnail=thumb_dir
                )
        except Exception as e:
            logger.error("MJ绘画回调失败: [%s]", str(e))
            if record:
                await UserWalletModel.inc(
                    scene=WalletEnum.POINTS,
                    user_id=record.user_id,
                    change_type=WalletEnum.POINTS_INC_DRAWS_FAIL,
                    change_amount=record.points,
                    project="AI绘画",
                    remarks=str(e)
                )
        finally:
            engine: str = await UrlUtil.get_storage_engine()
            if engine != "local":
                if album_dir_path and os.path.exists(album_dir_path):
                    os.unlink(album_dir_path)
                if thumb_dir_path and os.path.exists(thumb_dir_path):
                    os.unlink(thumb_dir_path)
